package tv.dyndns.kishibe.server.relevance;

import java.io.IOException;
import java.io.Reader;
import java.util.ArrayList;
import java.util.List;

import org.apache.lucene.analysis.Tokenizer;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.OffsetAttribute;
import org.apache.lucene.util.AttributeSource;

import tv.dyndns.kishibe.server.util.Normalizer;

import com.google.common.base.Throwables;
import com.google.common.io.CharStreams;

public class ViterbiTokenizer extends Tokenizer {
	private final CharTermAttribute termAtt;
	private final OffsetAttribute offsetAtt;
	private final List<Integer> offsets = new ArrayList<Integer>();
	private final List<Integer> lengths = new ArrayList<Integer>();
	private char[] buffer;
	private int wordIndex = 0;

	public ViterbiTokenizer(Reader input) {
		super(input);
		offsetAtt = addAttribute(OffsetAttribute.class);
		termAtt = addAttribute(CharTermAttribute.class);
		initialize();
	}

	public ViterbiTokenizer(AttributeFactory factory, Reader input) {
		super(factory, input);
		offsetAtt = addAttribute(OffsetAttribute.class);
		termAtt = addAttribute(CharTermAttribute.class);
		initialize();
	}

	public ViterbiTokenizer(AttributeSource source, Reader input) {
		super(source, input);
		offsetAtt = addAttribute(OffsetAttribute.class);
		termAtt = addAttribute(CharTermAttribute.class);
		initialize();
	}

	private void initialize() {
		offsets.clear();
		lengths.clear();

		String s = null;
		try {
			s = CharStreams.toString(input);
		} catch (Exception e) {
			throw Throwables.propagate(e);
		}
		s = Normalizer.normalize(s);
		Trie trie = null;
		try {
			trie = Trie.getInstance();
		} catch (Exception e) {
			e.printStackTrace();
		}
		trie.parse(s, null, offsets, lengths);
		buffer = s.toCharArray();
		wordIndex = 0;
	}

	@Override
	public boolean incrementToken() throws IOException {
		if (wordIndex == offsets.size()) {
			return false;
		} else {
			final int offset = offsets.get(wordIndex);
			final int length = lengths.get(wordIndex);
			++wordIndex;
			termAtt.copyBuffer(buffer, offset, length);
			offsetAtt.setOffset(offset, offset + length);
			return true;
		}
	}

	@Override
	public void end() throws IOException {
		super.end();
		offsetAtt.setOffset(buffer.length, buffer.length);
	}

	@Override
	public void reset(Reader input) throws IOException {
		super.reset(input);
		initialize();
	}
}
